package org.example.code.test2024.eleven;

import org.example.code.test2024.eight_mouth.TreeNode;

import java.util.*;

public class BinaryTreeShift {

    public TreeNode shiftTree(TreeNode root, int k) {
        if (root == null) {
            return null;
        }

        // 使用队列进行层次遍历
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);

        while (!queue.isEmpty()) {
            int size = queue.size();
            List<TreeNode> levelNodes = new ArrayList<>();

            // 收集当前层的所有节点
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                levelNodes.add(node);
                if (node.left != null) {
                    queue.offer(node.left);
                }
                if (node.right != null) {
                    queue.offer(node.right);
                }
            }

            // 提取当前层的节点值
            List<Integer> values = new ArrayList<>();
            for (TreeNode node : levelNodes) {
                values.add(node.val);
            }

            // 对当前层的值右移
            int n = values.size();
            int shift = k % n; // 实际右移位数
            Collections.rotate(values, shift);

            // 更新节点的值
            for (int i = 0; i < levelNodes.size(); i++) {
                levelNodes.get(i).val = values.get(i);
            }
        }

        return root;
    }

    // Helper method to print tree in level-order for verification
    public void printTree(TreeNode root) {
        if (root == null) {
            System.out.println("Tree is empty.");
            return;
        }
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                if (node != null) {
                    System.out.print(node.val + " ");
                    queue.offer(node.left);
                    queue.offer(node.right);
                } else {
                    System.out.print("null ");
                }
            }
            System.out.println();
        }
    }

    public static void main(String[] args) {
        BinaryTreeShift solution = new BinaryTreeShift();

        // 构造示例二叉树
        TreeNode root = new TreeNode(1);
        root.left = new TreeNode(2);
        root.right = new TreeNode(3);
        root.left.left = new TreeNode(4);
        root.left.right = new TreeNode(5);
        root.right.left = new TreeNode(6);
        root.right.right = new TreeNode(7);

        System.out.println("原始二叉树:");
        solution.printTree(root);

        // 右移每层节点
        int k = 2;
        TreeNode shiftedRoot = solution.shiftTree(root, k);

        System.out.println("右移后的二叉树:");
        solution.printTree(shiftedRoot);
    }
}
